/*
 * Copyright (C) 2016-2023 T-Head Semiconductor Co., Ltd. All rights reserved.
 *
 * SPDX-License-Identifier: Apache-2.0
 *
 * Licensed under the Apache License, Version 2.0 (the License); you may
 * not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an AS IS BASIS, WITHOUT
 * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "shl_gref.h"

int shl_gref_conv1d(struct csinn_tensor *input, struct csinn_tensor *output,
                    struct csinn_tensor *kernel, struct csinn_tensor *bias,
                    struct csinn_conv1d_params *params)
{
    shl_gref_sidcso_op(input, output, kernel, bias, CSINN_OP_CONV1D, params);
    return CSINN_TRUE;
}

int shl_gref_conv1d_infer_shape(struct csinn_tensor *input, struct csinn_tensor *output,
                                struct csinn_tensor *kernel, struct csinn_tensor *bias,
                                struct csinn_conv1d_params *params)
{
    int c, w, kernel_oc;
    shl_tensor_try_nc1xc0_to_ndarray_shape(input);
    if (input->layout == CSINN_LAYOUT_NCW) {
        c = 1;
        w = 2;
        kernel_oc = kernel->dim[0];
    } else if (input->layout == CSINN_LAYOUT_NWC) {
        w = 1;
        c = 2;
        kernel_oc = params->group == input->dim[c] ? kernel->dim[2] : kernel->dim[0];
    } else {
        shl_debug_error("%s: Invalid input tensor layout!\n", __func__);
        return CSINN_UNSUPPORT_LAYOUT;
    }

    int32_t in_w = input->dim[w];
    int32_t kernel_w = kernel->dim[w];
    int32_t padding_w = params->pad_left + params->pad_right;
    int32_t stride_w = params->stride_width;
    int32_t dilation_w = params->dilation_width;

    output->layout = input->layout;
    output->dim_count = 3;
    output->dim[0] = input->dim[0];  // N
    output->dim[c] = kernel_oc;
    output->dim[w] = (in_w + padding_w - dilation_w * (kernel_w - 1) - 1) / stride_w + 1;
    return CSINN_TRUE;
}

int shl_gref_depthwise_conv1d(struct csinn_tensor *input, struct csinn_tensor *output,
                              struct csinn_tensor *kernel, struct csinn_tensor *bias,
                              struct csinn_conv1d_params *params)
{
    shl_gref_sidcso_op(input, output, kernel, bias, CSINN_OP_DEPTHWISE_CONV1D, params);
    return CSINN_TRUE;
}

int shl_gref_depthwise_conv1d_infer_shape(struct csinn_tensor *input, struct csinn_tensor *output,
                                          struct csinn_tensor *kernel, struct csinn_tensor *bias,
                                          struct csinn_conv1d_params *params)
{
    return shl_gref_conv1d_infer_shape(input, output, kernel, bias, params);
}
